import csv

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import numpy as np
from datetime import datetime, timedelta

pains = ['无','轻微','中等','严重']
stresss = ['低','中','高']

# 数据预处理类
class PeriodDataset(Dataset):
    def __init__(self, date_infos, seq_length=3):
        """
        date_infos: 包含多维特征的数据矩阵 [样本数, 特征数]
        seq_length: 输入序列长度
        """
        # 转换为numpy数组
        self.data = np.array(date_infos, dtype=np.float32)
        
        # 构建序列样本
        self.samples = []
        for i in range(len(self.data) - seq_length):
            x = self.data[i:i+seq_length]  # 取连续seq_length天的数据 [seq_length, features]
            y = self.data[i+seq_length]     # 下一天的数据 [features]
            self.samples.append((x, y))
            
        # 归一化参数（每个特征单独归一化）
        self.mean = np.mean(self.data, axis=0)
        self.std = np.std(self.data, axis=0)
        self.std[self.std == 0] = 1e-6  # 避免除零错误

    def __len__(self):
        return len(self.samples)
    
    def __getitem__(self, idx):
        x, y = self.samples[idx]
        # 对每个特征进行归一化
        x = (x - self.mean) / self.std
        y = (y - self.mean) / self.std
        return torch.FloatTensor(x), torch.FloatTensor(y)

# LSTM模型定义
class PeriodPredictor(nn.Module):
    def __init__(self, input_size=4,  # 对应4个特征
                 hidden_size=32, 
                 num_layers=2,
                 output_size=4):  # 预测4个特征
        super().__init__()
        self.lstm = nn.LSTM(
            input_size=input_size,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True
        )
        self.fc = nn.Linear(hidden_size, output_size)
        
    def forward(self, x):
        # x shape: (batch_size, seq_length, input_size)
        out, _ = self.lstm(x)
        out = self.fc(out[:, -1, :])  # 取最后一个时间步
        return out

# 示例使用
if __name__ == "__main__":
    with open('test.csv', 'r', newline='', encoding='utf-8') as f:
        reader = csv.reader(f)
        i = 0
        date_infos = []
        for row in reader:
            if i > 0:
                date_infos.append([
                    (datetime.strptime(row[0], '%Y-%m-%d') - datetime(2000,1,1)).days,
                    int(row[1]),
                    pains.index(row[2]),
                    stresss.index(row[3]),
                ])
            i += 1
        print(date_infos)

    #示例数据（实际应替换为真实数据）


    # 准备数据集
    dataset = PeriodDataset(date_infos, seq_length=3)
    dataloader = DataLoader(dataset, batch_size=2, shuffle=True)

    # 初始化模型
    model = PeriodPredictor()
    criterion = nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

    # 训练循环
    for epoch in range(100):
        for x, y in dataloader:
            optimizer.zero_grad()
            outputs = model(x)

            loss = criterion(outputs, y)
            loss.backward()
            optimizer.step()
        if epoch % 10 == 0:
            print(f'Epoch {epoch}, Loss: {loss.item():.4f}')


    m = dataset.mean
    s = dataset.std
    # 预测示例修改
    with torch.no_grad():
        # 获取最后3天的数据作为测试输入
        test_data = date_infos[-3:]  # 取最后3个时间步
        test_input = (np.array(test_data) - dataset.mean) / dataset.std
        test_input = torch.FloatTensor(test_input).unsqueeze(0)  # 添加batch维度 → [1, 3, 4]
        
        prediction = model(test_input)
        # 反归一化
        denorm_pred = prediction.numpy() * dataset.std + dataset.mean
        
        # 解析预测结果
        next_date = datetime(2000,1,1) + timedelta(days=int(denorm_pred[0][0]))
        print(f"预测日期: {next_date.strftime('%Y-%m-%d')}")
        print(f"是否经期: {int(round(denorm_pred[0][1]))}")
        print(f"疼痛程度: {pains[int(round(denorm_pred[0][2]))]}")
        print(f"压力水平: {stresss[int(round(denorm_pred[0][3]))]}")